#include "BVP.h"
#include <algorithm>
#include <queue>
#include <iterator>
#include <cmath>
#include <lapacke.h>

//---------- PUBLIC INTERFACE --------------

BVP2d::BVP2d() : D{{0, 0}, 0}, f(nullptr) {}

BVP2d::BVP2d(Point2d center, double radius, const std::function<double(Point2d)> &f, BoundaryCondition BC)
    : f(f), BC(BC)
{
    D = Disk{center, radius};
}

BVP2d::~BVP2d()
{
}

void BVP2d::Discretization(int n)
{
    grid.clear();
    adjacancy.clear();

    double h = 1.0 / n;
    // add all the internal points inside the domain
    for (int i = 1; i < n; i++)
    {
        for (int j = 1; j < n; j++)
        {
            Point2d p = {i * h, j * h};
            if (PointStatus(p) == 0) // internal point
                grid.push_back({i * h, j * h});
        }
    }

    // gurantee there are at least 4 grid points
    if (grid.size() <= 4)
        throw std::domain_error("less than 4 grid points");

    // Link neighboring or boundary points
    int inter_num = grid.size();
    adjacancy.resize(inter_num);
    for (int i = 0; i < inter_num; ++i)
    {
        LinkSingle(i, h, 0);
        LinkSingle(i, -h, 0);
        LinkSingle(i, h, 1);
        LinkSingle(i, -h, 1);
    }

    // check connectness
    if (!isConnected(n, inter_num))
        throw std::domain_error("Use a connected domain or switch to a denser grid!");
}

std::vector<double> BVP2d::Solve(int n)
{
    // Discretize the domain into grids
    Discretization(n);

    double h = 1 / n;

    // the order of the matrix we need to solve
    int N = grid.size();

    double *A = new double[N * N]{0};
    double *b = new double[N]{0};
    int *IPIV = new int[N]{0}, *nrhs = new int(1), *INFO = new int;

    for (int i = 0; i < N; ++i)
    {
        const int status = PointStatus(grid[i]);
        if (status == -1)
        {
            throw std::domain_error("The grid has a point outside the domain!");
        }
        else if (status == 0) // internal points
        {
            // use FD method to discretize the equation -u_xx-u_yy = f;
            double rx = grid[adjacancy[i][0]][0] - grid[i][0];
            double lx = grid[i][0] - grid[adjacancy[i][1]][0];
            double uy = grid[adjacancy[i][2]][1] - grid[i][1];
            double dy = grid[i][1] - grid[adjacancy[i][3]][1];

            double row_coef = rx * lx * (rx + lx) / 2;
            double col_coef = uy * dy * (uy + dy) / 2;

            A[i + i * N] = (rx + lx) / row_coef + (uy + dy) / col_coef;
            A[i + adjacancy[i][0] * N] = -lx / row_coef;
            A[i + adjacancy[i][1] * N] = -rx / row_coef;
            A[i + adjacancy[i][2] * N] = -dy / col_coef;
            A[i + adjacancy[i][3] * N] = -uy / col_coef;

            b[i] = f(grid[i]);
        }
        else if (status == 1) // points on the square boundary
        {
            // get the index of the only internal point
            // that a boundary point connects
            int closest_idx = adjacancy[i][0];

            /**
             *        2
             *        |
             *   1  <-p->  0
             *        |
             *        3
             */
            // j is used to signify which direction the boundary point located
            // with respect to the internal point
            int j = 0;
            for (; j < 4; ++j)
            {
                if (adjacancy[closest_idx][j] == i)
                    break;
            }

            // get the opposite point the internal point connects along the same axis
            //  0<->1 and 2<->3 , which switches j to the opposite direction
            j = 2 * (j / 2) + (1 - j % 2);
            int closest_next = adjacancy[closest_idx][j];

            // represents the normal vector of the boundary
            int direction = (j == 0 || j == 2) ? 1 : -1;

            double dx1, dx2;

            int axis = j / 2;
            dx1 = grid[closest_idx][axis] - grid[i][axis];
            dx2 = grid[closest_next][axis] - grid[i][axis];

            A[i + i * N] = BC.a + direction * BC.b * (-(dx1 + dx2) / (dx1 * dx2));
            A[i + closest_idx * N] = direction * BC.b * (dx2 / (dx1 * (dx2 - dx1)));
            A[i + closest_next * N] = direction * BC.b * (-dx1 / (dx2 * (dx2 - dx1)));

            b[i] = BC.g(grid[i]);
        }
        else if (status == 2) // points on the circle
        {
            // normal vector
            Point2d direction = normalize(grid[i] - D.center);

            // the same as the previous operation
            int closest_idx = adjacancy[i][0];
            int j = 0;
            for (; j < 4; ++j)
            {
                if (adjacancy[closest_idx][j] == i)
                    break;
            }
            j = 2 * (j / 2) + (1 - j % 2);
            int closest_next = adjacancy[closest_idx][j];

            double dx1, dx2, component[2];

            // component[0] represents the component of normal vector along the axis j locates
            // component[1] represents the other component of normal vector
            int axis = j / 2;
            dx1 = grid[closest_idx][axis] - grid[i][axis];
            dx2 = grid[closest_next][axis] - grid[i][axis];
            component[0] = direction[axis];

            A[i + i * N] = BC.a + component[0] * BC.b * (-(dx1 + dx2) / (dx1 * dx2));
            A[i + closest_idx * N] = component[0] * BC.b * (dx2 / (dx1 * (dx2 - dx1)));
            A[i + closest_next * N] = component[0] * BC.b * (-dx1 / (dx2 * (dx2 - dx1)));

            // get the other two points along the other axis
            int otherside_idx1 = adjacancy[closest_idx][2 * (1 - axis)],
                otherside_idx2 = adjacancy[closest_idx][2 * (1 - axis) + 1];

            double dy1, dy2;

            dy1 = grid[otherside_idx1][1 - axis] - grid[closest_idx][1 - axis];
            dy2 = grid[closest_idx][1 - axis] - grid[otherside_idx2][1 - axis];
            component[1] = direction[1 - axis];

            double otherside_coef = dy1 * dy2 * (dy2 + dy1);

            A[i + closest_idx * N] += component[1] * BC.b * (dy1 * dy1 - dy2 * dy2) / otherside_coef;
            A[i + otherside_idx1 * N] += component[1] * BC.b * dy2 * dy2 / otherside_coef;
            A[i + otherside_idx2 * N] += component[1] * BC.b * (-dy1 * dy1) / otherside_coef;

            b[i] = BC.g(grid[i]);
        }
    }

    // for (int i = 0; i < N; ++i)
    // {
    //     for (int j = 0; j < N; ++j)
    //         std::cout << A[i + j * N] << " ";
    //     std::cout << std::endl;
    // }

    dgesv_(&N, nrhs, A, &N, IPIV, b, &N, INFO);

    std::vector<double> solution;
    if (*INFO == 0)
    {
        solution = std::vector<double>{b, b + N};
    }
    else if (*INFO > 0)
        std::cout << " lapack_dgesv error: A is singular " << std::endl;
    else
        std::cout << " lapack_dgesv: the " << -(*INFO) << "-th argument has an illegal value" << std::endl;

    delete[] A;
    delete[] b;
    delete nrhs;
    delete INFO;

    return solution;
}

std::vector<Point2d> &BVP2d::getGrid() { return grid; }

void BVP2d::PrintAdajacancy() const
{
    for (int i = 0; i < adjacancy.size(); ++i)
    {
        std::cout << grid[i] << ": ";
        for (int j = 0; j < adjacancy[i].size(); ++j)
        {
            std::cout << grid[adjacancy[i][j]] << ", ";
        }
        std::cout << std::endl;
    }
}

//---------- PRIVATE METHOD --------------

int BVP2d::PointStatus(Point2d p) const
{
    double epsilon = 1e-12;
    if (max_of(p) <= 1 && min_of(p) >= 0)
    {

        if (D.radius > 0)
        {
            Point2d vec = p - D.center;
            double dis = norm(vec, 2);
            if (fabs(dis - D.radius) < epsilon)
                return 2;
            else if (dis < D.radius)
                return -1;
        }
        if (max_of(p) == 1 || min_of(p) == 0)
        {
            return 1;
        }
        return 0;
    }

    return -1;
}

inline void BVP2d::LinkSingle(int i, double h, int axis)
{
    Point2d p = grid[i];
    // the point neiboring to p
    Point2d p1(p + Point2d{h * (1 - axis), h * axis});

    int status = PointStatus(p1);
    if (status >= 0)
    {
        auto it = std::find(grid.begin(), grid.end(), p1);

        if (it == grid.end()) // a new point
        {
            adjacancy[i].push_back(grid.size());
            grid.push_back(p1);
            adjacancy.push_back({i});
        }
        else // a point already exists
        {
            adjacancy[i].push_back(distance(grid.begin(), it));
        }
    }
    else
    {
        if (D.radius > 0 && max_of(p1) < 1 && min_of(p) > 0) // inside the disk
        {
            // add the point on the circle instead
            if (axis == 0)
            {
                double dx = sqrt(pow(D.radius, 2) - pow(p[1] - D.center[1], 2));
                (p[0] > D.center[0]) ? grid.push_back(Point2d{D.center[0] + dx, p[1]})
                                     : grid.push_back(Point2d{D.center[0] - dx, p[1]});
            }
            else if (axis == 1)
            {
                double dy = sqrt(pow(D.radius, 2) - pow(p[0] - D.center[0], 2));
                (p[1] > D.center[1]) ? grid.push_back(Point2d{p[0], D.center[1] + dy})
                                     : grid.push_back(Point2d{p[0], D.center[1] - dy});
            }
            adjacancy[i].push_back(grid.size() - 1);
            adjacancy.push_back({i});
        }
    }
}

bool BVP2d::isConnected(int n, int inter_num) const
{

    double h = 1 / n;

    // Connected_Component
    std::vector<int> CC;
    std::queue<int> q;

    CC.push_back(0);
    q.push(0);

    while (!q.empty())
    {
        int tmp = q.front();
        q.pop();
        for (int i = 0; i < 4; ++i)
        {
            int idx = adjacancy[tmp][i];
            if (idx < inter_num && std::find(CC.begin(), CC.end(), idx) == CC.end())
            {
                q.push(idx);
                CC.push_back(idx);
            }
        }
    }

    // check if all the internal points form a connected component
    return CC.size() == inter_num;
}
